from llama_index.llms.qianfan import Qianfan
import asyncio,qianfan
from llama_index.core import Settings
from llama_index.core.query_engine import NLSQLTableQueryEngine,SQLTableRetrieverQueryEngine
from sqlalchemy import (
    create_engine,
)
access_key = "access_key"
secret_key = "secret_key"
model_name = "ERNIE-Speed-128K"
endpoint_url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/ernie_speed"
context_window = 30000
from llama_index.core import Settings,get_response_synthesizer,SQLDatabase
from llama_index.core.agent import ReActAgent
from llama_index.core.tools import QueryEngineTool
from llama_index.embeddings.jinaai import JinaEmbedding
from llama_index.core.objects import (
SQLTableNodeMapping,
ObjectIndex,
SQLTableSchema,
)
from llama_index.core import (
    SimpleDirectoryReader,
    VectorStoreIndex,
    StorageContext,
    load_index_from_storage,
    Prompt
)
Settings.llm = Qianfan(access_key, secret_key, model_name, endpoint_url, context_window)

Settings.embed_model = JinaEmbedding(
    api_key="api_key",
    model="jina-embeddings-v2-base-en",
    embed_batch_size=16,
)

engine = create_engine("sqlite:///data/test.db")
sql_database = SQLDatabase(engine,include_tables=["students","grade"],sample_rows_in_table_info=3,view_support=True) 
table_node_mapping = SQLTableNodeMapping(sql_database)
table_schema_objs = [
    (SQLTableSchema(table_name="students",context_str=("crew表是学生表,存储的是Name名字(text),age职位(int),address住址(str),sex性别(str)"))),
    (SQLTableSchema(table_name="grade",context_str=("ports表是分数表,存储的是students_id学生表id(int),english英语成绩(int),maths数学成绩(int),chinese语文成绩(int)"))),
]
obj_index = ObjectIndex.from_objects(
    table_schema_objs,
    table_node_mapping,
    VectorStoreIndex,
)
query_engine = SQLTableRetrieverQueryEngine(
    sql_database= sql_database, 
    table_retriever= obj_index.as_retriever(similarity_top_k=2),
)
response = query_engine.query('数学最高学生的信息,他的各科成绩怎么样？')
print(response)

# agent = ReActAgent.from_tools([sql_tool],verbose=True)
# response = agent.chat('数学最高分是谁？')
# print(response)


